import  torch
from torch import nn, optim
from torch.nn import functional as F
import torchvision
from    torch.utils.data import DataLoader
from    torchvision import datasets
from    torchvision import transforms

rnn = nn.RNN(input_size=100, hidden_size=20, num_layers=1)
print(rnn)
x = torch.randn(10, 3, 100)
out, h = rnn(x, torch.zeros(1, 3, 20))
print(out.shape, h.shape)

rnncell1 = nn.RNNCell(100, 30)
rnncell2 = nn.RNNCell(30, 20)
h1 = torch.zeros(3, 30)
h2 = torch.zeros(3, 20)

for xt in x.split(1, dim=0):
    xt = xt.squeeze(0)
    h1 = rnncell1(xt, h1)
    h2 = rnncell2(h1, h2)

print(h2.shape)

lstm = nn.LSTM(input_size=100, hidden_size=20, num_layers=4) # 四层LSTM
print(lstm)
x = torch.randn(10, 128, 100)
out, (h, c) = lstm(x)
print(out.shape, h.shape, c.shape)